package com.sunzm.spark.sql.udaf

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator
import org.roaringbitmap.RoaringBitmap

/**
 * 将一组 序列化的 BitMap 字节数组反序列化后 进行合并, 然后再进行序列化返回
 */
class BitMapMergeUDAF extends Aggregator[Array[Byte], Array[Byte],Array[Byte]]{
  override def zero: Array[Byte] = {
    // 构造一个空的bitmap
    val bm: RoaringBitmap = RoaringBitmap.bitmapOf()

    // 序列化这个空的bitmap交还给框架
    BitMapUtil.serBitMap(bm)
  }

  override def reduce(buffer1: Array[Byte], buffer2: Array[Byte]): Array[Byte] = {

    BitMapUtil.mergeBitmapByteArray(buffer1, buffer2)

  }

  override def merge(buffer1: Array[Byte], buffer2: Array[Byte]): Array[Byte] = {
    BitMapUtil.mergeBitmapByteArray(buffer1, buffer2)
  }

  override def finish(reduction: Array[Byte]): Array[Byte] = {
    reduction
  }

  override def bufferEncoder: Encoder[Array[Byte]] = Encoders.BINARY

  override def outputEncoder: Encoder[Array[Byte]] = Encoders.BINARY
}
